import sqlite3
import pandas as pd
from sqlalchemy import create_engine, Column, Integer, String, Float, Date, DateTime, Boolean, Text, MetaData, Table
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import os

# 数据库配置
DATABASE_URL = "sqlite:///./data/stock_data.db"

# 确保数据目录存在
os.makedirs("./data", exist_ok=True)

# SQLAlchemy 配置
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# 股票财务数据表模型
class StockFinancial(Base):
    __tablename__ = "stock_financial"

    id = Column(Integer, primary_key=True, index=True)
    stock_code = Column(String, index=True)  # 股票代码
    stock_name = Column(String)  # 股票名称
    year = Column(Integer, index=True)  # 年份
    revenue = Column(Float)  # 营业收入
    net_profit = Column(Float)  # 净利润
    report_date = Column(Date)  # 报告日期

# 用户表模型
class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True, nullable=False)
    email = Column(String, unique=True, index=True, nullable=False)
    full_name = Column(String)
    hashed_password = Column(String, nullable=False)
    is_active = Column(Boolean, default=True)
    is_admin = Column(Boolean, default=False)
    avatar_url = Column(String)
    bio = Column(Text)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_login = Column(DateTime)

def create_tables():
    """创建数据库表"""
    Base.metadata.create_all(bind=engine)

def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def save_financial_data_to_db(df, stock_code, stock_name):
    """将财务数据保存到数据库"""
    db = SessionLocal()
    try:
        # 清除该股票的旧数据
        db.query(StockFinancial).filter(StockFinancial.stock_code == stock_code).delete()
        
        # 插入新数据
        for _, row in df.iterrows():
            financial_record = StockFinancial(
                stock_code=stock_code,
                stock_name=stock_name,
                year=int(row.get('年份', 0)),
                revenue=float(row.get('营业收入', 0)) if pd.notna(row.get('营业收入', 0)) else 0,
                net_profit=float(row.get('净利润', 0)) if pd.notna(row.get('净利润', 0)) else 0,
                report_date=pd.to_datetime(row.get('报告日期', '1900-01-01')).date()
            )
            db.add(financial_record)
        
        db.commit()
        return True
    except Exception as e:
        db.rollback()
        print(f"保存数据时出错: {e}")
        return False
    finally:
        db.close()
